"""
RAG Embeddings - 向量查询示例模块

该模块演示如何使用向量数据库进行文档检索：
- 连接到已有的向量数据库
- 执行相似度查询
- 展示查询结果

Author: xiaodeme@163.com
Date: 2025-07-09 18:10:00
"""

from vector_db import VectorDB


class QueryExample:
    """向量查询示例类"""
    
    def __init__(self):
        """初始化查询器，连接向量数据库"""
        print("[INFO] 初始化向量数据库连接...")
        self.vector_db = VectorDB()
        print("[INFO] 向量数据库连接成功")

    def single_query(self, query_text, n_results=3):
        """
        执行单个查询并显示结果
        
        Args:
            query_text (str): 查询文本
            n_results (int): 返回结果数量，默认3个
            
        Returns:
            list: 查询到的文档列表
        """
        print(f"\n{'='*50}")
        print(f"查询: {query_text}")
        print(f"{'='*50}")
        
        # 执行查询
        results = self.vector_db.query(query_text, n_results=n_results)
        
        # 提取文档内容
        documents_data = results.get('documents', [[]])
        documents = documents_data[0] if documents_data and len(documents_data) > 0 else []
        
        if documents:
            print(f"\n找到 {len(documents)} 个相关文档:")
            for i, doc in enumerate(documents, 1):
                print(f"\n📄 文档 {i}:")
                print(f"内容: {doc}")
                print("-" * 50)
        else:
            print("未找到相关文档")
            
        return documents

    def batch_query(self, query_list, n_results=3):
        """
        批量执行查询
        
        Args:
            query_list (list): 查询文本列表
            n_results (int): 每个查询返回的结果数量
        """
        print(f"\n🔍 开始批量查询，共 {len(query_list)} 个查询...")
        
        for i, query in enumerate(query_list, 1):
            print(f"\n🔸 查询 {i}/{len(query_list)}")
            self.single_query(query, n_results)


def main():
    """主函数：演示向量数据库查询功能"""
    
    try:
        # 初始化查询器
        query_example = QueryExample()
        
        # 定义测试查询列表
        test_queries = [
            "市场准入可以包括哪些方面?",
            "创新器械特别审批重点关注哪些方面?", 
            "什么时候搭建质量管理体系?",
            "软件类产品如何选择检测机构?",
            "临床试验大致可以分为几个阶段?"
        ]
        
        # 获取集合信息
        info = query_example.vector_db.get_collection_info()
        print(f"当前数据库包含 {info['document_count']} 个文档")
        
        if info['document_count'] == 0:
            print("\n⚠️  向量数据库为空，请先运行 vector_db.py 添加文档")
            return
        
        # 演示单个查询
        print("\n🎯 单个查询演示:")
        query_example.single_query(test_queries[0])
        
        # 询问是否继续批量查询
        print(f"\n是否执行所有 {len(test_queries)} 个测试查询？(y/n): ", end="")
        user_input = input().strip().lower()
        
        if user_input in ['y', 'yes', '是']:
            # 执行批量查询
            query_example.batch_query(test_queries)
        else:
            print("查询演示结束")
            
    except Exception as e:
        print(f"[ERROR] 查询过程中发生错误: {e}")
        print("请确保已经运行 vector_db.py 初始化数据库")


if __name__ == '__main__':
    main()